

#include "Symtab.h"
#include "Type.h"
#include "CFG.h"
#include "CodeSource.h"
#include "CodeObject.h"

#include "InstructionDecoder.h"
#include "Instruction.h"

#include "feature.h"
#include "graphlet.h"
#include "colors.h"
#include "supergraph.h"

#include "BPatch_binaryEdit.h"
#include "BPatch_image.h"
#include "BPatch_object.h"
#include "BPatch_function.h"
#include "BPatch_flowGraph.h"
#include "BPatch_basicBlock.h"
#include "BPatch_basicBlockLoop.h"

#include "config.h"
#include "feature_analyzer.h"
#include "types.h"
#include "statistics.h"

#include <cstdio>
#include <cstring>

#include <string>
#include <vector>
#include <map>
#include <set>

#include <unordered_map>
#include <unordered_set>

#include <boost/iterator/filter_iterator.hpp>

using boost::make_filter_iterator;

using namespace std;
using namespace Dyninst;
using namespace Dyninst::ParseAPI;
using namespace Dyninst::SymtabAPI;
using namespace Dyninst::InstructionAPI;
using namespace graphlets;

bool FeatureAnalyzer::NORT = true;
bool FeatureAnalyzer::NOPLT = true;
int FeatureAnalyzer::NODES = 3;
int FeatureAnalyzer::N = 3;
bool FeatureAnalyzer::ANON = false;
int FeatureAnalyzer::MERGE = 1;

BPatch FeatureAnalyzer::bpatch;


FeatureAnalyzer::FeatureAnalyzer(Config *c):
    config(c) {}

int FeatureAnalyzer::Analyze(char *filename) {
    
    BPatch_binaryEdit *openedFile = bpatch.openBinary(filename);
    if (openedFile == NULL) {
        fprintf(stderr, "cannot open binary file %s\n", filename);
	return -1;
    }       
    BPatch_image *image = openedFile->getImage();
    vector<BPatch_object*> objs;
    image->getObjects(objs);
    if (objs.size() != 1) {
        fprintf(stderr, "Unknown BPatch_objects\n");
	for (auto iter = objs.begin(); iter != objs.end(); ++iter)
	    printf("%s %s\n", (*iter)->name().c_str(), (*iter)->pathName().c_str());
	return -1;
    }

    Statistics stat;
    
    CodeObject *co = ParseAPI::convert(objs[0]);
    // Is it safe to convert from BPatch_object to Symtab, then create a new SymtabCodeSource
    SymtabCodeSource *sts = new SymtabCodeSource(SymtabAPI::convert(objs[0]));
            
    RecordAllInternalFunction(image, co);

    unordered_set<Address> visitedBlocks;

    vector<BPatch_function*> *funcs = image->getProcedures();
    
    string id(strrchr(filename, '/')+1);
    char buf[64];
    for(auto fit = funcs->begin(); fit != funcs->end(); ++fit){
        BPatch_function *bpatch_func = *fit;
        ParseAPI::Function * f = ParseAPI::convert(bpatch_func);
        
        if (strncmp(f->name().c_str(),"std::",5) == 0 || strncmp(f->name().c_str(),"__gnu_cxx::",11) == 0) {
            //fprintf(stderr,"skipped %s\n",f->name().c_str());
            continue;
        }
        if (!f->blocks().empty()) {
            if (NORT && f->src() == RT)
                continue;
            
            if (NOPLT && sts->linkage().find(f->addr()) != sts->linkage().end())
                continue;
            

            BPatch_flowGraph * cfg = bpatch_func->getCFG();

	    vector<BPatch_basicBlockLoop*> outerLoops;
	    cfg->getOuterLoops(outerLoops);

	    set<BPatch_basicBlock*> basicBlocks;
	    cfg->getAllBasicBlocks(basicBlocks);

	    set<BPatch_basicBlock*> loopBasicBlocks;
	    int index = 0;
	    for (auto loopIt = outerLoops.begin(); loopIt != outerLoops.end(); ++loopIt) {
	        if ((*loopIt)->parent != NULL) {
		    fprintf(stderr, "%s\n", "Warning: outer loops returned by cfg->getOuterLoops() still have parent!");
		    continue;
		}
	        vector<BPatch_basicBlock*> bb;
		(*loopIt)->getLoopBasicBlocks(bb);
		loopBasicBlocks.insert(bb.begin(), bb.end());
		++index;
		snprintf(buf, 64, "%d", index);
		AnalyzeALoop( (id + "+" + bpatch_func->getName() + "+outerloops+loop_" + string(buf)).c_str(), sts, f, *loopIt);
	    }

	    /*  Currently we only care about loops
	     * 
	    for (auto bit = basicBlocks.begin(); bit != basicBlocks.end(); ++bit){
	        if (loopBasicBlocks.find(*bit) != loopBasicBlocks.end()) {
		    // This basic block is in a loop. It has been taken care of
		    continue;
		}
	        ParseAPI::Block *b = ParseAPI::convert(*bit);

                // Functions may share blocks
                if (visitedBlocks.find(b->start()) != visitedBlocks.end())
		    continue;

		visitedBlocks.insert(b->start());
		stat.Increment();
		if (config->idiomFeature) ProduceIdiomFeature(f, b);
		if (config->ngramFeature) ProduceNGramFeature(sts, b);            
		if (config->graphletFeature) ProduceGraphletFeature(b);
		
		//if (config->libcallFeature) ProduceLibcallFeature(co, f);
		//if (config->dataFeature) ProduceDataFeature(f, obj);
                

		if (config->addressingModeFeature) ProduceAddressingModeFeature(b);
	    }
            */

        }
    }
    
    stat.Print(config->statisticFile, filename);

    totalStat.Accumulate(stat);
    delete sts;

    return 0;

}

int FeatureAnalyzer::AnalyzeALoop(const char* id, SymtabCodeSource *sts, ParseAPI::Function *f, BPatch_basicBlockLoop *loop) {
    if (config->idiomFeature) ProduceLoopIdiomFeature(id, f, loop);
    if (config->ngramFeature) ProduceLoopNGramFeature(id, sts, loop);
    if (config->graphletFeature) ProduceLoopGraphletFeature(id, f, loop);
    if (config->addressingModeFeature) ProduceLoopAddressingModeFeature(id, loop);

}


void FeatureAnalyzer::RecordAllInternalFunction(BPatch_image *image, CodeObject *co){
    internalFuncs.clear();
    vector<BPatch_function*> *funcs = image->getProcedures();
    for (auto fit = funcs->begin(); fit != funcs->end(); ++fit){
        BPatch_function *bpatch_func = *fit;
	ParseAPI::Function *parse_func = ParseAPI::convert(bpatch_func);
        if (co->cs()->linkage().find(parse_func->addr()) == co->cs()->linkage().end()){
	    vector<string> names;
	    bpatch_func->getMangledNames(names);
	    internalFuncs.insert(names.begin(), names.end());
        }
    }
}

void FeatureAnalyzer::PrintFeatureList() {
    if (!config->featureList) return;
    for (auto it = featureIndex.begin(); it != featureIndex.end(); ++it) {
        fprintf(config->featureListFile, "%d %s %s\n", it->second.first, it->first.c_str(), it->second.second.c_str());
    }
}

